﻿#define _WIN32_WINNT 0x0600
#include <windows.h>
#include "MinHook.h"
#include <sstream>
#include <iomanip>
#include <fstream>
#include <string>

#pragma comment(lib, "MinHook.x64.lib")

// 导出函数
extern "C" {
    __declspec(dllexport) BOOL InstallGlobalHooks();
    __declspec(dllexport) BOOL UninstallGlobalHooks();
    __declspec(dllexport) void SetMainProcessId(DWORD pid);
}

// 全局变量
static DWORD g_mainProcessId = 0;
static bool g_hooksInstalled = false;

// 原始函数指针
static HHOOK(WINAPI* TrueSetWindowsHookEx)(int, HOOKPROC, HINSTANCE, DWORD) = nullptr;
static BOOL(WINAPI* TrueSetForegroundWindow)(HWND) = nullptr;

// 日志函数
void LogToFile(const wchar_t* message) {
    std::wofstream logFile(L"C:\\HookLog.txt", std::ios::app);
    if (logFile) {
        logFile << message << std::endl;
        logFile.close();
    }
}

// 辅助函数：将HWND格式化为十六进制字符串
std::wstring HwndToString(HWND hWnd) {
    std::wstringstream ss;
    ss << L"0x" << std::hex << std::setw(16) << std::setfill(L'0') << reinterpret_cast<ULONG_PTR>(hWnd);
    return ss.str();
}

// 钩子函数
HHOOK WINAPI HookedSetWindowsHookEx(int idHook, HOOKPROC lpfn, HINSTANCE hMod, DWORD dwThreadId) {
    // 检查是否是鼠标/键盘钩子
    if (idHook == WH_MOUSE || idHook == WH_KEYBOARD ||
        idHook == WH_MOUSE_LL || idHook == WH_KEYBOARD_LL) {

        // 获取调用者进程ID
        DWORD callerPid = GetCurrentProcessId();

        // 如果不是主进程，阻止钩子安装
        if (callerPid != g_mainProcessId) {
            std::wstring logMsg = L"[GlobalHook] Blocked hook installation from non-main process. Caller PID: ";
            logMsg += std::to_wstring(callerPid);
            logMsg += L", Hook Type: ";
            logMsg += std::to_wstring(idHook);
            LogToFile(logMsg.c_str());
            return NULL;
        }
    }

    return TrueSetWindowsHookEx(idHook, lpfn, hMod, dwThreadId);
}

BOOL WINAPI HookedSetForegroundWindow(HWND hWnd) {
    // 检查是否是目标窗口
    HWND hHost = reinterpret_cast<HWND>(GetPropW(hWnd, L"HOST_WINDOW"));

    if (hHost && IsWindow(hHost)) {
        std::wstring logMsg = L"[GlobalHook] Blocked SetForegroundWindow for target window: ";
        logMsg += HwndToString(hWnd);
        logMsg += L", Host window: ";
        logMsg += HwndToString(hHost);
        LogToFile(logMsg.c_str());
        return TRUE;  // 完全阻止激活
    }
    LogToFile(L"[GlobalHook]gn");
    return TrueSetForegroundWindow(hWnd);
}

// DLL入口点
BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved) {
    switch (ul_reason_for_call) {
    case DLL_PROCESS_ATTACH:
        DisableThreadLibraryCalls(hModule);
        // 在新进程中安装钩子
        if (InstallGlobalHooks()) {
            std::wstring logMsg = L"[DLL] Hook installed in process: ";
            logMsg += std::to_wstring(GetCurrentProcessId());
            LogToFile(logMsg.c_str());
        }
        break;
    case DLL_PROCESS_DETACH:
        if (g_hooksInstalled) {
            UninstallGlobalHooks();
        }
        break;
    }
    return TRUE;
}

// 设置主进程ID
void SetMainProcessId(DWORD pid) {
    g_mainProcessId = pid;
    std::wstring logMsg = L"[GlobalHook] SetMainProcessId: ";
    logMsg += std::to_wstring(pid);
    LogToFile(logMsg.c_str());
}

// 安装全局钩子
BOOL InstallGlobalHooks() {
    if (g_hooksInstalled) {
        LogToFile(L"[GlobalHook] Hooks already installed");
        return TRUE;
    }

    LogToFile(L"[GlobalHook] Initializing MinHook");
    if (MH_Initialize() != MH_OK) {
        LogToFile(L"[GlobalHook] MH_Initialize failed");
        return FALSE;
    }

    // 钩住SetWindowsHookExA和W
    LogToFile(L"[GlobalHook] Creating hook for SetWindowsHookExA");
    if (MH_CreateHookApiEx(L"user32", "SetWindowsHookExA",
        reinterpret_cast<LPVOID>(HookedSetWindowsHookEx),
        reinterpret_cast<LPVOID*>(&TrueSetWindowsHookEx), nullptr) != MH_OK) {
        LogToFile(L"[GlobalHook] Failed to hook SetWindowsHookExA");
        return FALSE;
    }

    LogToFile(L"[GlobalHook] Creating hook for SetWindowsHookExW");
    if (MH_CreateHookApiEx(L"user32", "SetWindowsHookExW",
        reinterpret_cast<LPVOID>(HookedSetWindowsHookEx),
        reinterpret_cast<LPVOID*>(&TrueSetWindowsHookEx), nullptr) != MH_OK) {
        LogToFile(L"[GlobalHook] Failed to hook SetWindowsHookExW");
        return FALSE;
    }

    // 钩住SetForegroundWindow
    LogToFile(L"[GlobalHook] Creating hook for SetForegroundWindow");
    if (MH_CreateHookApiEx(L"user32", "SetForegroundWindow",
        reinterpret_cast<LPVOID>(HookedSetForegroundWindow),
        reinterpret_cast<LPVOID*>(&TrueSetForegroundWindow), nullptr) != MH_OK) {
        LogToFile(L"[GlobalHook] Failed to hook SetForegroundWindow");
        return FALSE;
    }

    LogToFile(L"[GlobalHook] Enabling hooks");
    if (MH_EnableHook(MH_ALL_HOOKS) != MH_OK) {
        LogToFile(L"[GlobalHook] MH_EnableHook failed");
        return FALSE;
    }

    g_hooksInstalled = true;
    LogToFile(L"[GlobalHook] Hooks installed successfully");
    return TRUE;
}

// 卸载全局钩子
BOOL UninstallGlobalHooks() {
    if (!g_hooksInstalled) {
        LogToFile(L"[GlobalHook] Hooks not installed");
        return TRUE;
    }

    LogToFile(L"[GlobalHook] Disabling hooks");
    MH_DisableHook(MH_ALL_HOOKS);

    LogToFile(L"[GlobalHook] Uninitializing MinHook");
    MH_Uninitialize();

    g_hooksInstalled = false;
    LogToFile(L"[GlobalHook] Hooks uninstalled");
    return TRUE;
}